package tmhprediction.eval;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;

import tmhprediction.classification.TMHClassifier;
import tmhprediction.main.TMHResultMap;


public class TMHResidueEvaluator {
	
	public static int[][] calculateConfusionMatrix(HashMap<String, double[]> result)
	{
		int[][] cm = new int[2][2];
    	int tp = 0;
    	int fp = 0;
    	int fn = 0;
    	int tn = 0;
		
		for(String residue : result.keySet())
		{
			double[] values = result.get(residue);
		    //Building confusion matrix
		    if(values[0] == TMHClassifier.PARSEDSYMBOLTMHRESIDUE) //obs
		    {
		    	if(values[1] == TMHClassifier.PARSEDSYMBOLTMHRESIDUE)
		    	{
		    		tp++;
		    	}
		    	else
		    	{
		    		fn++;
		    	}
		    }
		    else
		    {
		    	if(values[1] == TMHClassifier.PARSEDSYMBOLTMHRESIDUE)
		    	{
		    		fp++;
		    	}
		    	else
		    	{
		    		tn++;
		    	}
		    }
		}
		
		cm[0][0] = tp;
		cm[0][1] = fn;
		cm[1][0] = fp;
		cm[1][1] = tn;

		return cm;
	}
	
    public static void calcPerResidueScores(HashMap<String, double[]> resultMap, String outputPath) {
    
    int[][] confMatrix = TMHResidueEvaluator.calculateConfusionMatrix(resultMap);	
    	
	double tp = confMatrix[0][0];
	double fp = confMatrix[1][0];
	double fn = confMatrix[0][1];
	double tn = confMatrix[1][1];
	double qObs2T = tp / (tp + fn);
	double qPred2T = tp / (tp + fp);
	double qObs2N = tn / (tn + fp);
	double qPred2N = tn / (tn + fn);
	double q2 = (tp + tn) / (tp + tn + fp + fn);
	double mcc = ((tp * tn) - (fp * fn))
		/ Math.sqrt(((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)));
	double sens = tp / (tp + fn);
	double spec = tn / (tn + fp);
    	
//	MSE calculation
//	double rmse = eval.rootMeanSquaredError();	// TODO self-implement
	
	
	double mse = 0;
	for (String residue : resultMap.keySet())
	{
		double confidenceTMH = resultMap.get(residue)[2];
		double confidenceNonTMH =  1.0 - confidenceTMH;
		
		double oTm, oNonTm;
		double observed = resultMap.get(residue)[0];	
		if(observed == TMHClassifier.PARSEDSYMBOLTMHRESIDUE)
		{
			oTm = 1.0;
			oNonTm = 0.0;
		}
		else
		{
			oTm = 0.0;
			oNonTm = 1.0;
		}
		
		mse = mse + ( (Math.pow(oTm - confidenceTMH, 2) + Math.pow(oNonTm - confidenceNonTMH,2)) / 2.0 );	
	}
	
	mse = (1.0 / resultMap.keySet().size()) * mse;
	
	
	BufferedWriter writer = null;
	try {
	    File file = new File(outputPath);
	    if (!file.exists()) {
	    	(new File(file.getParent())).mkdirs();
	    	file.createNewFile();
	    }
	    writer = new BufferedWriter(new FileWriter(file));
	} catch (Exception ex) {
	    ex.printStackTrace();
	}
	try {
	    writer.append("TP:" + tp + "\n");
	    writer.append("TN: " + tn + "\n");
	    writer.append("FP: " + fp + "\n");
	    writer.append("FN: " + fn + "\n");
	    writer.append("Q^Obs_2T: " + qObs2T + "\n");
	    writer.append("Q^Pred_2T: " + qPred2T + "\n");
	    writer.append("Q^Obs_2N: " + qObs2N + "\n");
	    writer.append("Q^Pred_2N: " + qPred2N + "\n");
	    writer.append("Q2: " + q2 + "\n");
//	    writer.append("RMSE: " + rmse + "\n");
	    writer.append("MSE:" + mse + "\n");
	    writer.append("MCC: " + mcc + "\n");
	    writer.append("Sens: " + sens + "\n");
	    writer.append("Spec: " + spec + "\n");
	    writer.close();
	} catch (IOException e) {
	    // TODO Auto-generated catch block
	    e.printStackTrace();
	}
    }
    
    public static double[] calcPerResidueScoresSilent(HashMap<String, double[]> resultMap)
    {
        
        int[][] confMatrix = TMHResidueEvaluator.calculateConfusionMatrix(resultMap);	
        	
    	double tp = confMatrix[0][0];
    	double fp = confMatrix[1][0];
    	double fn = confMatrix[0][1];
    	double tn = confMatrix[1][1];
    	double qObs2T = tp / (tp + fn);
    	double qPred2T = tp / (tp + fp);
    	double qObs2N = tn / (tn + fp);
    	double qPred2N = tn / (tn + fn);
    	double q2 = (tp + tn) / (tp + tn + fp + fn);
    	double mcc = ((tp * tn) - (fp * fn))
    		/ Math.sqrt(((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)));
    	double sens = tp / (tp + fn);
    	double spec = tn / (tn + fp);
        	
//    	MSE calculation
//    	double rmse = eval.rootMeanSquaredError();	// TODO self-implement
    	
    	
    	double mse = 0;
    	for (String residue : resultMap.keySet())
    	{
    		double confidenceTMH = resultMap.get(residue)[2];
    		double confidenceNonTMH =  1.0 - confidenceTMH;
    		
    		double oTm, oNonTm;
    		double observed = resultMap.get(residue)[0];	
    		if(observed == TMHClassifier.PARSEDSYMBOLTMHRESIDUE)
    		{
    			oTm = 1.0;
    			oNonTm = 0.0;
    		}
    		else
    		{
    			oTm = 0.0;
    			oNonTm = 1.0;
    		}
    		
    		mse = mse + ( (Math.pow(oTm - confidenceTMH, 2) + Math.pow(oNonTm - confidenceNonTMH,2)) / 2.0 );	
    	}
    	
    	mse = (1.0 / resultMap.keySet().size()) * mse;
    	
    	double[] results = new double[3];
    	results[0] = q2;  
    	results[1] = qPred2T; 
    	results[2] = qObs2T;
    	
    	return results;
    			
    }

    
    public static TMHResultMap evaluateProteins(TMHResultMap map) {
    	TMHResultMap result = new TMHResultMap();
        
        for (String idPos : map.keySet()) {
		    String parsedName = parseProteinNameFromIdPos(idPos);
		    double[] mapValues = map.get(idPos);
		    double[] values = new double[3];
		    if (result.get(parsedName) == null) {
			values[0] = mapValues[0];
			values[1] = mapValues[1];
			values[2] = 0.0;
			
			if (mapValues[1] == 1.0) {
			    values[2] = mapValues[2];
			}
		    } else {
				double[] oldValues = result.get(parsedName);
				values[0] = mapValues[0] + oldValues[0];
				values[1] = mapValues[1] + oldValues[1];
				values[2] = oldValues[2];
				if (mapValues[1] == 1.0) {
				    values[2] += mapValues[2];
				}
		}
	    result.put(parsedName, values.clone());
	}
	return result;
    }

    private static String parseProteinNameFromIdPos(String idPos) {
	String result = "";
	String[] splitString = idPos.split("_");
	for (int i = 0; i < splitString.length - 1; i++) {
	    if (i != 0) {
		result += "_";
	    }
	    result += splitString[i];
	}
	return result;
    }
    
    public static void generateQ2ListFile(TMHResultMap map, String outputFilePath) throws IOException {
	File file = new File(outputFilePath);
	if (!file.exists()) {
	    file.getParentFile().mkdirs();
	    file.createNewFile();
	}
	BufferedWriter writer = new BufferedWriter(new FileWriter(file));
	HashMap<String, Integer> allProtMap = new HashMap<String, Integer>();
	for (String name : map.keySet()) {
	    allProtMap.put(parseProteinNameFromIdPos(name), 1);
	}
	for (String protName : allProtMap.keySet()) {
	    double[] eval = TMHResidueEvaluator.calcPerResidueScoresSilent(getProteinMap(protName, map));
	    writer.append(protName + "\t" + Double.toString(eval[0]));
	    writer.newLine();
	}
	writer.close();
    }
    
    public static TMHResultMap getProteinMap(String proteinName, TMHResultMap map) {
	TMHResultMap result = new TMHResultMap();
	for (String name : map.keySet()) {
	    if (parseProteinNameFromIdPos(name).equals(proteinName)) {
		result.put(name, map.get(name));
	    }
	}
	return result;
    }
    
    public static void generateHelixFiles(TMHResultMap map, String observedOut, String predictedOut) throws Exception {
	File obs = new File(observedOut);
	if (!obs.exists()) {
	    obs.getParentFile().mkdirs();
	    obs.createNewFile();
	}
	File pred = new File(predictedOut);
	if (!pred.exists()) {
	    pred.getParentFile().mkdirs();
	    pred.createNewFile();
	}
	BufferedWriter obsWriter = new BufferedWriter(new FileWriter(obs));
	BufferedWriter predWriter = new BufferedWriter(new FileWriter(pred));
	HashMap<String, Integer> allProtMap = new HashMap<String, Integer>();
	for (String name : map.keySet()) {
	    allProtMap.put(parseProteinNameFromIdPos(name), 1);
	}
	for (String protName : allProtMap.keySet()) {
	    boolean obsLastIsHelix = false;
	    boolean predLastIsHelix = false;
	    int obsCurrentHelixCounter = 0;
	    int predCurrentHelixCounter = 0;
	    int obsHelixes = 0;
	    int predHelixes = 0;
	    TMHResultMap protMap = getProteinMap(protName, map);
	    ArrayList<String> list = new ArrayList<String>(protMap.keySet());
	    Collections.sort(list, new Comparator<String>() {

		@Override
		public int compare(String arg0, String arg1) {
		    int a = parseNumberFromName(arg0);
		    int b = parseNumberFromName(arg1);
		    return Integer.valueOf(a).compareTo(Integer.valueOf(b));
		}
		
		private int parseNumberFromName(String name) {
		    String[] field = name.split("_");
		    return Integer.parseInt(field[field.length - 1]);
		}
	    });
	    for (String res : list) {
		if (protMap.getObserved(res) == TMHClassifier.PARSEDSYMBOLTMHRESIDUE) {
		    if (!obsLastIsHelix) {
			obsLastIsHelix = true;
			obsHelixes += 1;
		    }
		    obsCurrentHelixCounter += 1;
		} else {
		    if (obsLastIsHelix) {
			obsLastIsHelix = false;
			obsWriter.append(protName + "_" + obsHelixes + "\t" + obsCurrentHelixCounter);
			obsWriter.newLine();
			obsCurrentHelixCounter = 0;
		    }
		}
		if (protMap.getPredicted(res) == TMHClassifier.PARSEDSYMBOLTMHRESIDUE) {
		    if (!predLastIsHelix) {
			predLastIsHelix = true;
			predHelixes += 1;
		    }
		    predCurrentHelixCounter += 1;
		} else {
		    if (predLastIsHelix) {
			predLastIsHelix = false;
			predWriter.append(protName + "_" + predHelixes + "\t" + predCurrentHelixCounter);
			predWriter.newLine();
			predCurrentHelixCounter = 0;
		    }
		}
	    }
	}
	obsWriter.close();
	predWriter.close();
    }
}
